#!/usr/bin/env python3
import rospy
import subprocess
import sys
import os
import yaml
from sensor_msgs.msg import Image
from cv_bridge import CvBridge

class RTMPStreamer:
    def __init__(self):
        rospy.init_node('rgb2rtmp_streamer')
        self.bridge = CvBridge()
        self.process_dict = {}  # topic -> ffmpeg进程
        self.size_dict = {}     # topic -> (w, h)
        self.topic_url_map = {} # topic -> rtmp_url
        self.subscribers = []

        config_path = os.path.join(os.path.dirname(__file__), '../config.yaml')
        try:
            if os.path.exists(config_path):
                with open(config_path, 'r') as f:
                    config = yaml.safe_load(f)
                if isinstance(config, dict):
                    base_url = config.get('base_url')
                    streams = config.get('streams', [])
                else:
                    base_url = None
                    streams = []
                if not base_url or not streams:
                    rospy.logerr("config.yaml 缺少 base_url 或 streams 字段，无法启动推流")
                    sys.exit(1)
                for stream in streams:
                    topic = stream.get('topic')
                    suffix = stream.get('suffix')
                    if not topic or not suffix:
                        rospy.logwarn(f"配置项缺少topic或suffix: {stream}")
                        continue
                    rtmp_url = base_url + suffix
                    self.topic_url_map[topic] = rtmp_url
                    sub = rospy.Subscriber(topic, Image, self.image_callback, callback_args=topic)
                    self.subscribers.append(sub)
                    rospy.loginfo(f"已为话题 {topic} 配置推流: {rtmp_url}")
                if not self.topic_url_map:
                    rospy.logerr("config.yaml 中未找到有效的 streams 配置，无法启动推流")
                    sys.exit(1)
            else:
                rospy.logerr("未找到 config.yaml，无法启动推流")
                sys.exit(1)
        except Exception as e:
            rospy.logerr(f"加载 config.yaml 失败: {str(e)}，无法启动推流")
            sys.exit(1)

        rospy.on_shutdown(self.cleanup)

    def image_callback(self, msg, topic):
        try:
            cv_image = self.bridge.imgmsg_to_cv2(msg, "bgr8")
            if topic not in self.process_dict:
                # 首帧，记录尺寸并启动ffmpeg
                h, w = cv_image.shape[:2]
                self.size_dict[topic] = (w, h)
                self.start_ffmpeg(topic, w, h)
            proc = self.process_dict.get(topic)
            if proc and proc.stdin:
                proc.stdin.write(cv_image.tobytes())
        except Exception as e:
            rospy.logerr(f"[{topic}] 图像处理错误: {str(e)}")

    def start_ffmpeg(self, topic, width, height):
        rtmp_url = self.topic_url_map[topic]
        ffmpeg_cmd = [
            'ffmpeg',
            '-f', 'rawvideo',
            '-pix_fmt', 'bgr24',
            '-s', f'{width}x{height}',
            '-r', '30',
            '-fflags', 'nobuffer',
            '-avioflags', 'direct',
            '-i', '-',
            '-c:v', 'libx264',
            '-pix_fmt', 'yuv420p',
            '-preset', 'ultrafast',
            '-tune', 'zerolatency',
            '-f', 'flv',
            rtmp_url
        ]
        try:
            proc = subprocess.Popen(
                ffmpeg_cmd,
                stdin=subprocess.PIPE
            )
            self.process_dict[topic] = proc
            rospy.loginfo(f"[{topic}] 启动ffmpeg推流到: {rtmp_url}")
        except Exception as e:
            rospy.logerr(f"[{topic}] 启动ffmpeg失败: {str(e)}")
            self.process_dict[topic] = None

    def cleanup(self):
        for topic, proc in self.process_dict.items():
            if proc:
                try:
                    if proc.stdin:
                        proc.stdin.close()
                    proc.terminate()
                    proc.wait(timeout=2)
                    rospy.loginfo(f"[{topic}] ffmpeg进程已停止")
                except Exception as e:
                    rospy.logerr(f"[{topic}] 停止ffmpeg时出错: {str(e)}")

if __name__ == '__main__':
    try:
        streamer = RTMPStreamer()
        rospy.spin()
    except rospy.ROSInterruptException:
        pass